#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 1; // tensor num for each queue
constexpr int32_t CALC_ALIGN_NUM =256;

class KernelGcd {
public:
    __aicore__ inline KernelGcd() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, 
                                uint32_t bigDataCoreNum, uint32_t smallBlockLength, uint32_t bigBlockLength,
                                uint32_t smallTileNum, uint32_t smallTileLength, uint32_t smallLasttileLength,
                                uint32_t bigTileNum, uint32_t bigTileLength, uint32_t bigLasttileLength)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        
        // AscendC::printf("this->times = %d \n", this->times);
        
        if(AscendC::GetBlockIdx() >= bigDataCoreNum)
        {
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1 + (bigDataCoreNum * bigBlockLength) + (smallBlockLength * (AscendC::GetBlockIdx() - bigDataCoreNum)), smallBlockLength* this->times);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2 + (bigDataCoreNum * bigBlockLength) + (smallBlockLength * (AscendC::GetBlockIdx() - bigDataCoreNum)), smallBlockLength* this->times);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + (bigDataCoreNum * bigBlockLength) + (smallBlockLength * (AscendC::GetBlockIdx() - bigDataCoreNum)), smallBlockLength * this->times);
            pipe.InitBuffer(inQueueX1, BUFFER_NUM, smallTileLength * sizeof(DTYPE_X1));
            pipe.InitBuffer(inQueueX2, BUFFER_NUM, smallTileLength * sizeof(DTYPE_X1));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, smallTileLength * sizeof(DTYPE_X1));
            this->tileNum = smallTileNum;
            this->tileLength = smallTileLength;
            this->lasttileLength = smallLasttileLength;         
        }
        else
        {
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1 + bigBlockLength * AscendC::GetBlockIdx(), bigBlockLength);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x2 + bigBlockLength * AscendC::GetBlockIdx(), bigBlockLength);
            yGm.SetGlobalBuffer((__gm__ DTYPE_X1*)y + bigBlockLength * AscendC::GetBlockIdx(), bigBlockLength);
            pipe.InitBuffer(inQueueX1, BUFFER_NUM, bigTileLength * sizeof(DTYPE_X1));
            pipe.InitBuffer(inQueueX2, BUFFER_NUM, bigTileLength * sizeof(DTYPE_X1));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, bigTileLength * sizeof(DTYPE_X1));
            this->tileNum = bigTileNum;
            this->tileLength = bigTileLength;
            this->lasttileLength = bigLasttileLength;  
        }  

        
        // pipe.InitBuffer(tempBuf, (this->tileLength * sizeof(int32_t)));
        // this->tempLocal = tempBuf.Get<int32_t>();
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum;
        this->calcLength = this->tileLength;
        for (int32_t i = 0; i < loopCount; i++) {
            if(i == loopCount -1)
            {
                this->calcLength = this->lasttileLength;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        AscendC::LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        AscendC::LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X1>();
        AscendC::DataCopy(x1Local, x1Gm[progress * this->tileLength], this->calcLength);
        AscendC::DataCopy(x2Local, x2Gm[progress * this->tileLength], this->calcLength);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }

    __aicore__ inline void Compute(int32_t progress)
    {
        AscendC::LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        AscendC::LocalTensor<DTYPE_X1> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        AscendC::LocalTensor<DTYPE_X1> yLocal = outQueueY.AllocTensor<DTYPE_X1>();
        
        for (int i=0; i<(this->calcLength); i++) {
            DTYPE_X1 x1 = x1Local.GetValue(i);
            DTYPE_X2 x2 = x2Local.GetValue(i);
            x1 = abs(x1);
            x2 = abs(x2);
            while (x2 != 0) {
                DTYPE_X2 temp = x2;
                x2 = x1 % x2;
                x1 = temp;
            }
            yLocal.SetValue(i, x1);
        }

        outQueueY.EnQue<DTYPE_X1>(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        AscendC::LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        AscendC::DataCopy(yGm[progress * this->tileLength], yLocal, this->calcLength);
        outQueueY.FreeTensor(yLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    AscendC::TBuf<AscendC::TPosition::VECCALC>  tempBuf;

    AscendC::GlobalTensor<DTYPE_X1> x1Gm;
    AscendC::GlobalTensor<DTYPE_X2> x2Gm;
    AscendC::GlobalTensor<DTYPE_Y> yGm;

    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t calcLength;
    uint32_t lasttileLength;
    uint32_t times;
};

class KernelGcd_Broadcast {
public:
    __aicore__ inline KernelGcd_Broadcast() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t totalLength,
                                uint32_t dim,
                                uint32_t* shape, uint32_t* reduce1, uint32_t* reduce2
                                ) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->reduce1 = reduce1;
        this->reduce2 = reduce2;
        this->shape = shape;
        this->dim = dim;
        this->totalLength = totalLength;

        this->tileLength = 4096;

        // get start index for current core, core parallel
        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x2);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X1*)y);

        // AscendC::printf("x1Gm_addr = %f \n", x1Gm.GetPhyAddr());
        // AscendC::printf("x2Gm_addr = %f \n", x2Gm.GetPhyAddr());

        pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    }

    __aicore__ inline void Process() {
        int32_t count = this->totalLength / this->shape[this->dim - 1];
        this->lastDimLength = this->shape[this->dim - 1];
        uint32_t elemOfBlock = 32 / sizeof(DTYPE_X1);
        this->lastDimLengthAlign = (this->lastDimLength + elemOfBlock -1) / elemOfBlock * elemOfBlock;

        // 假定最后一维的数据可以完全搬入ub
        // uint32_t totalLength = this->shape[this->dim - 1];
        // this->tileNum = totalLength / this->tileLength + (totalLength % this->tileLength > 0);
        uint32_t d[21] = {0};
        uint32_t dn1[21] = {0};
        uint32_t dn2[21] = {0};
        uint32_t dn3[21] = {0};

        auto dim = this->dim - 1;
        d[dim] = dn1[dim] = dn2[dim] = dn3[dim] = 1;
        InitializeDnArrays(d, dn1, this->reduce1, dim, this->shape);
        InitializeDnArrays(d, dn2, this->reduce2, dim, this->shape);

        

        for(int j=0; j<count; j++){
            uint32_t start1 = 0, start2 = 0;
            CalculateStart(j, start1, dn1, reduce1, d);
            CalculateStart(j, start2, dn2, reduce2, d);
            // AscendC::printf("j=%d, start1 = %d, start2 = %d" , j, start1, start2);
            CopyIn(start1 * this->lastDimLength, start2 * this->lastDimLength, this->lastDimLengthAlign);
            Compute();
            CopyOut(j * this->lastDimLength, this->lastDimLengthAlign);
        }   

    }

private:
    __aicore__ inline void CopyIn(uint32_t start1, uint32_t start2, uint32_t length) {//both 1，n->m,n
        AscendC::LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        AscendC::LocalTensor<DTYPE_X1> x2Local = inQueueX2.AllocTensor<DTYPE_X1>();

        AscendC::DataCopy(x1Local, x1Gm[start1], length);
        AscendC::DataCopy(x2Local, x2Gm[start2], length);



        // AscendC::printf("CopyIn start1 = %d", start1);
        // AscendC::printf("CopyIn start2 = %d", start2);

        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    __aicore__ inline void Compute() {
        AscendC::LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        AscendC::LocalTensor<DTYPE_X1> x2Local = inQueueX2.DeQue<DTYPE_X1>();

        AscendC::LocalTensor<DTYPE_X1> yLocal = outQueueY.AllocTensor<DTYPE_X1>();

        // AscendC::printf("%d");
        // AscendC::DumpTensor(x1Local, 1001, (this->lastDimLengthAlign));
        // AscendC::printf("%d");
        // AscendC::DumpTensor(x2Local, 1002, (this->lastDimLengthAlign));

        for (int i=0; i<(this->lastDimLength); i++) {
            DTYPE_X1 x1 = x1Local.GetValue(i);
            DTYPE_X2 x2 = x2Local.GetValue(i);
            x1 = abs(x1);
            x2 = abs(x2);
            while (x2 != 0) {
                DTYPE_X2 temp = x2;
                x2 = x1 % x2;
                x1 = temp;
            }
            yLocal.SetValue(i, x1);
        }


        outQueueY.EnQue<DTYPE_X1>(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }

    __aicore__ inline void CopyOut(uint32_t start, uint32_t length) {
        AscendC::LocalTensor<DTYPE_X1> yLocal = outQueueY.DeQue<DTYPE_X1>();
        AscendC::DataCopy(yGm[start], yLocal, length);
        outQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void CalculateStart(int j, uint32_t &start, uint32_t* dn, uint32_t* reduce, uint32_t* d) {
        for (int k = dim - 1; k >= 0; k--) {
            uint32_t index = (j / d[k + 1] % shape[k]);

            if (reduce[k] == 0) {
                start += dn[k + 1] * index;
            }
        }
    }
    __aicore__ inline void InitializeDnArrays(uint32_t* d, uint32_t* dn, uint32_t* reduce, uint32_t dim, uint32_t* shape) {
        d[dim] = dn[dim] = 1;
        for (int k = dim - 1; k >= 0; k--) {
            d[k] = d[k + 1] * shape[k];
            if (reduce[k] == 0) {
                dn[k] = dn[k + 1] * shape[k];
            } else {
                dn[k] = dn[k + 1];
            }
        }
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueY;


    AscendC::GlobalTensor<DTYPE_X1> x1Gm;
    AscendC::GlobalTensor<DTYPE_X1> x2Gm;
    AscendC::GlobalTensor<DTYPE_X1> yGm;

    uint32_t blockLength;
    uint32_t startLength;
    uint32_t startPointer;
    uint32_t weightLength;
    uint32_t endLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t mode;
    uint32_t position;

    uint32_t totalLength;
    uint32_t lastDimLengthAlign;
    uint32_t lastDimLength;

    uint32_t* reduce1;
    uint32_t* reduce2;
    uint32_t* shape;
    uint32_t dim;
};





extern "C" __global__ __aicore__ void gcd(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(TILING_KEY_IS(1)){
        KernelGcd op1;
        op1.Init(x1, x2, y, 
                tiling_data.bigDataCoreNum, tiling_data.smallBlockLength, tiling_data.bigBlockLength,
                tiling_data.smallTileNum, tiling_data.smallTileLength, tiling_data.smallLasttileLength, 
                tiling_data.bigTileNum, tiling_data.bigTileLength, tiling_data.bigLasttileLength);
        op1.Process();  
    }
    else if(TILING_KEY_IS(2))
    {
        KernelGcd_Broadcast op2;
        op2.Init(x1, x2, y, 
                tiling_data.totalLength,
                tiling_data.dim,
                tiling_data.shape, tiling_data.reduce1, tiling_data.reduce2);
        op2.Process();          
    }      

}